#!/usr/bin/env python2

#Author: Jamie Davis (davisjam@vt.edu)
#Description: Package for transforming callback descriptions into a graph
# Defines the following public classes: 
#	 Callback
#  CallbackGraph
# Python version: 2.7.6

import logging, argparse
import Callback as CB
import Palette
import gv # http://www.graphviz.org/pdf/gv.3python.pdf, http://www.graphviz.org/doc/info/attrs.html
import sys
import re

logger = logging.getLogger('root')
LOG_FORMAT = "[%(filename)s:%(lineno)s - %(funcName)20s() ] %(message)s"
logging.basicConfig(level=logging.DEBUG, format=LOG_FORMAT)
sys.setrecursionlimit(20000)

#input: (node, graph)
#output: nodeHandle corresponding to node
#	Adds node to graph if not yet present
def getNodeHandle (node, graph):
	# Registration ID is unique; execution ID is not since all unexecuted nodes have ID -1
	nodeName = "Node {}".format(node.getRegID())
	nodeHandle = gv.findnode(graph, nodeName)
	if (not nodeHandle):
		nodeHandle = gv.node(graph, "Node {}".format(node.getRegID()))
	return nodeHandle

#add NODE, its children, and a parental edge to GRAPH
#node handles for NODE and its children will be created if they do not already exist in GRAPH
#ACTION nodes will be ellipses. RESPONSE nodes will be squares.
def addNodeAndChildren (node, graph):
	assert(isinstance(node, CB.CallbackNode))
	
	nodeHandle = getNodeHandle(node, graph)
	
	#assign attributes
	extraInfo = node.getExtraInfo()
	if (extraInfo):
		extraInfo = "\\nExtra info: {}".format(extraInfo)
	else:
		extraInfo = ""
	nodeStr = "node <Reg {}, Exec {}>\\nTree level {}, Tree entry {}\\nCB type {}\\nContext {}\\nBehavior {}\\nExecuted {} {}".format(node.getRegID(), node.getExecID(), node.getTreeLevel(), node.getLevelEntry(), node.getCBType(), node.getContext(), node.getBehavior(), node.executed(), extraInfo)
	gv.setv(nodeHandle, "label", nodeStr)
	
	#assign shape based on behavior
	behavior = node.getBehavior()
	if (behavior == "ACTION" or behavior == "RESPONSE"):
		# All user code is put in squares so that the 'striped' style works
		shape = "square"
	else:
		shape = "ellipse"
	gv.setv(nodeHandle, "shape", shape)
	
	#Edges: parent -> children
	for child in node.children:
		childHandle = getNodeHandle(child, graph)
		gv.edge(nodeHandle, childHandle)
	#Edges: node dependencies -> node
	for dep in node.dependencies:
		depHandle = getNodeHandle(dep, graph)
		edgeHandle = gv.edge(depHandle, nodeHandle)
		gv.setv(edgeHandle, "style", "dotted")

#Returns true if node was not executed
#For use with CallbackTree.removeNodes
def nodeNotExecuted (node):
	assert(isinstance(node, CB.CallbackNode))
	return (not node.executed())

# input: (t, unit)
# t: time in nanoseconds
# unit: convert to this unit. one of 's', 'ms', 'us', 'ns'
# Returns "startTime" expressed in UNIT with no fractional component
def getTimeStr (t, unit='ns'):
	unitDivisors = {
		'ns': 1,
		'us': 1e3,
		'ms': 1e6,
		's':  1e9,
	}
	if (unit not in unitDivisors):
		assert (not "cbn_getTimeRangeStr: Error, invalid unit '{}'; select from {}".format(unit, unitDivisors.keys()))
	divisor = unitDivisors[unit]
	t = round(int(t)/divisor)
	return "{:.0f}".format(t)

#returns Node information in the following format:
# id:regTime:startTime:endTime
# times in usec
def cbn_adjListStr (node):	
 	return "{}:{}:{}:{}".format(node.getID(),
															getTimeStr(int(node.getRegistrationTime()), unit='us'),
															getTimeStr(int(node.getStartTime()), unit='us'),
															getTimeStr(int(node.getEndTime()), unit='us'))
	
#tree-walk function
#prints adjacency list information to f for each executed node
def walk_adjacencyList (node, f):
	if (not node.executed()):
		return
	myStr = cbn_adjListStr(node)
	childStrs = [cbn_adjListStr(n) for n in node.getChildren()]
	f.write("{} {}\n".	format(myStr, ",".join(childStrs)))

#Returns true if node is a marker event rather than a user CB
#For use with CallbackTree.removeNodes
def nodeIsMarker (node):
	assert(isinstance(node, CB.CallbackNode))
	return node.isMarkerNode()

#Returns true if node is non-user code rather than a user CB
#For use with CallbackTree.removeNodes
def nodeIsNonUserCode (node):
	assert(isinstance(node, CB.CallbackNode))
	return (not node.isUserCode())

def addExecOrder (tree, digraph):	
	cbNodes = tree.getExecOrder()
	for i in range(0, len(cbNodes)-1):
		j = i + 1
		node1, node2 = cbNodes[i], cbNodes[j]
		#Don't draw edges for un-executed nodes
		if (node1.executed() and node2.executed()):
			h1, h2 = getNodeHandle(node1, digraph), getNodeHandle(node2, digraph)
			edgeHandle = gv.edge(h1, h2)
			gv.setv(edgeHandle, "style", "invis")

# input: (rgb) tuple (R,G,B)
# output: (colorStr) a color string for graphviz
def rgbToGVColor (rgb):
	return "#%2x%2x%2x" % (rgb[0], rgb[1], rgb[2])

def addColors(tree, digraph, coloredNodes):
	allColoredNodes = {}  # Avoid duplicate entries; maps ID to node
	palette = Palette.Palette(len(coloredNodes))

	treeExecOrder = tree.getExecOrder()
	for colorGroup in coloredNodes:
		logging.info("colorGroup {}".format(colorGroup))
		colorGroupNodes = [n for n in treeExecOrder if int(n.getExecID()) in colorGroup]
		logging.info("colorGroup {} colorGroupNodes {} -- should be same length".format(colorGroup, colorGroupNodes))
		assert(len(colorGroup) == len(colorGroupNodes))

		rgb = palette.nextColor()

		for n in colorGroupNodes:
			colors = getattr(n, "_colors", [])
			colors.append(rgb)
			setattr(n, "_colors", colors)
			logging.info("node {} colors {}".format(n.getExecID(), colors))
			allColoredNodes[n.getExecID()] = n

	for n in allColoredNodes.values():
		h = getNodeHandle(n, digraph)
		logging.info("Coloring node {} (handle {})".format(n.getExecID(), h))

		kv = {}

		rgbs = n._colors
		fillColor = ':'.join([rgbToGVColor(rgb) for rgb in rgbs])
		kv["fillcolor"] = fillColor

		if len(n._colors) == 1:
			kv["style"] = "filled"
		else:
			shape = gv.getv(h, "shape")
			if shape == "ellipse":
				kv["style"] = "wedged"
			elif shape == "square":
				kv["style"] = "striped"
			else:
				raise Error("Error, unexpected shape {}".format(shape))

		for k in kv.keys():
			gv.setv(h, k, kv[k])
		del n._colors

def main():
	parser = argparse.ArgumentParser(description="Turn a libuv event schedule into a graph in .gv and/or adjacency list format")
	parser.add_argument("--schedFile", help="file containing libuv event schedule", required=True, type=str)

	parser.add_argument("--noMarkers", help="do not include the \"marker\" nodes that indicate uv loop progress in the schedule",	action="store_true")
	parser.add_argument("--onlyUserCode", help="do not include CBs for non-user code",	action="store_true")

	parser.add_argument("--gv", help="generate gv output (provide --gvF)", action="store_true")
	parser.add_argument("--gvF", help="gv output file", type=str)
	parser.add_argument("--gvOnlyExecuted", help="only include nodes that were executed", action="store_true")
	parser.add_argument("--gvExecOrder", help="vertical position indicates relative execution order in the gv graph", action="store_true")
	parser.add_argument("--gvColorFile", help="color the nodes specified in this file (format: one ID per line. If multiple colors, add lines like 'color i' for integer 0 <= N)", type=str)

	parser.add_argument("--adj", help="generate adjacency list output (provide --adjF). includes only executed nodes. ID is regID.", action="store_true")
	parser.add_argument("--adjF", help="adjacency list output file", type=str)

	args = parser.parse_args()

	logging.info("main: schedFile {}".format(args.schedFile))
	tree = CB.CallbackNodeTree(args.schedFile)

	if (args.noMarkers):
		logging.info("Removing marker nodes")
		tree.removeNodes(nodeIsMarker)

	if (args.onlyUserCode):
		logging.info("Removing any non-user code")
		tree.removeNodes(nodeIsNonUserCode)

	if (args.gv):
		assert(args.gvF)
		logging.info("graphviz: Generating graphviz version of schedule")

		coloredNodes = [] # list of lists of grouped nodes
		if (args.gvColorFile):
			logging.info("graphviz: Reading gvColorFile {}".format(args.gvColorFile))
			try:
				coloredNodes = CB.CallbackNodeGroups(args.gvColorFile).getNodeGroups()
			except IOError:
				logging.error("graphviz: Reading gvColorFile {} failed".format(args.gvColorFile))
				raise
			
		if (args.gvOnlyExecuted):
			logging.info("graphviz: Removing unexecuted nodes")
			tree.removeNodes(nodeNotExecuted)

		digraph = gv.digraph("graphviz: Creating gv graph from tree")
		tree.walk(addNodeAndChildren, digraph)

		if (args.gvExecOrder):
			logging.info("graphviz: Tweaking graph so it displays in execution order")
			addExecOrder(tree, digraph)

		if (args.gvColorFile):
			logging.info("graphviz: Adding colors")
			addColors(tree, digraph, coloredNodes)
	
		gv.write(digraph, args.gvF)
		logging.info("graphviz: Examine {} for the graphviz format".format(args.gvF))

	if (args.adj):
		assert(args.adjF)
		logging.info("adj: Generating adjacency list")

		try:
			with open(args.adjF, 'w') as f:
				tree.walk(walk_adjacencyList, f)
		except IOError:
			logging.error("adj: Writing to {} failed".format(args.adjF))
		logging.info("adj: Examine {} for the adjacency list".format(args.adjF))


###################################

main()
